"""
Event emission decorators for non-intrusive event integration.

This module provides decorator functions that enable easy integration of event
emission into existing training code without requiring major refactoring. The
decorators follow the Aspect-Oriented Programming (AOP) pattern, allowing
cross-cutting concerns like event emission to be added declaratively.

Key Features:
- Non-intrusive event integration
- Automatic context extraction from method arguments
- Error handling without disrupting training flow
- Conditional event emission based on custom criteria
- Performance-aware batch and epoch event handling

Decorator Types:
- emit_events: General-purpose decorator for before/after event emission
- emit_event_before/after: Simplified decorators for single event emission
- emit_on_condition: Conditional event emission based on custom logic
- emit_high_loss_event: Specialized decorator for loss monitoring
- emit_batch_events/emit_epoch_events: Training lifecycle decorators

Example:
    >>> from learners.events.decorators import emit_events, emit_high_loss_event
    >>> from learners.events.types import EventType
    >>>
    >>> class Trainer:
    ...     @emit_events(
    ...         before=EventType.ON_TRAIN_EPOCH_START,
    ...         after=EventType.ON_TRAIN_EPOCH_END,
    ...         extract_args={'epoch': 'epoch_num'}
    ...     )
    ...     def train_epoch(self, epoch_num, data_loader):
    ...         # Training logic here
    ...         pass
    ...
    ...     @emit_high_loss_event(threshold_key="loss_threshold")
    ...     def compute_loss(self, outputs, targets):
    ...         loss = criterion(outputs, targets)
    ...         return loss

Architecture:
    The decorators use introspection to extract method arguments and create
    appropriate event contexts. They're designed to be composable and can
    be combined with other decorators. Error handling ensures that event
    emission failures don't affect the training process.
"""

from functools import wraps
from typing import Callable, Optional, Any, Dict
import logging

from learners.interfaces.event import EventType


logger = logging.getLogger(__name__)


def emit_events(
    before: Optional[EventType] = None,
    after: Optional[EventType] = None,
    on_error: Optional[EventType] = None,
    condition: Optional[Callable] = None,
    extract_args: Optional[Dict[str, str]] = None,
) -> Callable:
    """
    Decorator to emit events before and/or after method execution.

    This is the most flexible decorator, allowing emission of different events
    at different points in method execution with full control over context
    extraction and conditional emission.

    Args:
        before: Event to emit before method execution
        after: Event to emit after method execution
        on_error: Event to emit if method raises an exception
        condition: Optional function to determine if events should be emitted
        extract_args: Mapping of event context keys to method argument names

    Example:
        @emit_events(
            before=EventType.ON_TRAIN_EPOCH_START,
            after=EventType.ON_TRAIN_EPOCH_END,
            extract_args={'epoch': 'epoch_num', 'task': 'task_id'}
        )
        def train_epoch(self, epoch_num, task_id, data_loader):
            # Training logic here
            pass

    Returns:
        Decorated function with event emission capability
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            # Check if instance has event emission capability
            if not hasattr(self, "emit_event"):
                return func(self, *args, **kwargs)

            # Check condition if provided
            if condition and not condition(self, *args, **kwargs):
                return func(self, *args, **kwargs)

            # Extract context data from method arguments
            event_context = {}
            if extract_args:
                event_context.update(
                    _extract_context_from_args(extract_args, func, args, kwargs)
                )

            # Add method metadata
            event_context.update(
                {
                    "method_name": func.__name__,
                    "method_args": args,
                    "method_kwargs": kwargs,
                }
            )

            # Emit before event
            if before:
                self.emit_event(before, **event_context)

            try:
                # Execute the original method
                result = func(self, *args, **kwargs)

                # Add result to context for after event
                if after:
                    after_context = event_context.copy()
                    after_context["result"] = result
                    self.emit_event(after, **after_context)

                return result

            except Exception as e:
                # Emit error event if specified
                if on_error:
                    error_context = event_context.copy()
                    error_context.update({"error": e, "error_type": type(e).__name__})
                    self.emit_event(on_error, **error_context)

                # Re-raise the exception
                raise

        return wrapper

    return decorator


def emit_event_before(event_type: EventType, **event_kwargs) -> Callable:
    """
    Decorator to emit an event before method execution.

    Simplified decorator for cases where only a before event is needed.

    Args:
        event_type: Type of event to emit before method execution
        **event_kwargs: Additional context to include with the event

    Returns:
        Decorator function
    """
    return emit_events(before=event_type, extract_args=event_kwargs)


def emit_event_after(event_type: EventType, **event_kwargs) -> Callable:
    """
    Decorator to emit an event after method execution.

    NOTE: This is essentially a simplified wrapper around emit_events().
    Could be considered redundant, but provides convenience.
    """
    return emit_events(after=event_type, extract_args=event_kwargs)


def emit_on_condition(
    condition: Callable, event_type: EventType, **event_kwargs
) -> Callable:
    """
    Decorator to emit an event only when a condition is met.

    NOTE: This duplicates some logic from emit_events() condition parameter.
    Could potentially be refactored to reduce code duplication.
    """

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            result = func(self, *args, **kwargs)

            if hasattr(self, "emit_event") and condition(self, result, *args, **kwargs):
                context = event_kwargs.copy()
                context.update({"method_name": func.__name__, "result": result})
                self.emit_event(event_type, **context)

            return result

        return wrapper

    return decorator


def emit_high_loss_event(
    threshold_key: str = "high_loss_threshold", default_threshold: float = 10.0
) -> Callable:
    """Decorator specifically for detecting and emitting high loss events."""

    def condition(self, result, *args, **kwargs):
        if hasattr(result, "item"):  # PyTorch tensor
            loss_value = result.item()
        elif isinstance(result, (int, float)):
            loss_value = result
        else:
            return False

        threshold = getattr(self, "args", {}).get(threshold_key, default_threshold)
        return loss_value > threshold

    return emit_on_condition(
        condition=condition, event_type=EventType.ON_HIGH_LOSS_DETECTED
    )


def emit_batch_events(
    start_event: EventType = EventType.ON_TRAIN_BATCH_START,
    end_event: EventType = EventType.ON_TRAIN_BATCH_END,
    extract_batch_idx: bool = True,
) -> Callable:
    """Decorator for batch-level event emission with automatic batch index extraction."""

    def decorator(func: Callable) -> Callable:
        @wraps(func)
        def wrapper(self, *args, **kwargs):
            if not hasattr(self, "emit_event"):
                return func(self, *args, **kwargs)

            # Try to extract batch_idx from arguments
            context: Dict[str, Any] = {"method_name": func.__name__}
            if extract_batch_idx:
                batch_idx = _find_batch_idx_in_args(func, args, kwargs)
                if batch_idx is not None:
                    context["batch_idx"] = batch_idx

            # Emit start event
            self.emit_event(start_event, **context)

            try:
                result = func(self, *args, **kwargs)

                # Emit end event with result
                end_context = context.copy()
                end_context["result"] = result
                self.emit_event(end_event, **end_context)

                return result
            except Exception as e:
                # Optionally emit error event
                error_context = context.copy()
                error_context["error"] = e
                logger.error(f"Error in {func.__name__}: {e}")
                raise

        return wrapper

    return decorator


def emit_epoch_events(
    start_event: EventType = EventType.ON_TRAIN_EPOCH_START,
    end_event: EventType = EventType.ON_TRAIN_EPOCH_END,
) -> Callable:
    """
    Decorator for epoch-level event emission.

    REDUNDANT: This is just a simple wrapper around emit_events().
    Consider if this convenience function is necessary.
    """
    return emit_events(before=start_event, after=end_event)


def emit_task_events(
    start_event: EventType = EventType.ON_TASK_START,
    end_event: EventType = EventType.ON_TASK_END,
) -> Callable:
    """
    Decorator for task-level event emission.

    REDUNDANT: Another simple wrapper around emit_events().
    Could be consolidated with other similar decorators.
    """
    return emit_events(before=start_event, after=end_event)


# Helper functions


def _extract_context_from_args(
    extract_args: Dict[str, str], func: Callable, args: tuple, kwargs: Dict[str, Any]
) -> Dict[str, Any]:
    """Extract context data from method arguments based on mapping."""
    context = {}

    # Get function signature for positional argument mapping
    import inspect

    sig = inspect.signature(func)
    param_names = list(sig.parameters.keys())[1:]  # Skip 'self'

    # Create argument mapping
    arg_dict = dict(zip(param_names, args))
    arg_dict.update(kwargs)

    # Extract specified arguments
    for context_key, arg_name in extract_args.items():
        if arg_name in arg_dict:
            context[context_key] = arg_dict[arg_name]

    return context


def _find_batch_idx_in_args(
    func: Callable, args: tuple, kwargs: Dict[str, Any]
) -> Optional[int]:
    """Try to find batch_idx in method arguments."""
    # Common parameter names for batch index
    batch_idx_names = ["batch_idx", "batch_index", "idx", "i"]

    # Check kwargs first
    for name in batch_idx_names:
        if name in kwargs:
            return kwargs[name]

    # Check positional args by function signature
    import inspect

    try:
        sig = inspect.signature(func)
        param_names = list(sig.parameters.keys())[1:]  # Skip 'self'

        for i, param_name in enumerate(param_names):
            if param_name in batch_idx_names and i < len(args):
                return args[i]
    except Exception:
        pass

    return None
